import torch
import torch.nn as nn
import torch.optim as optim

# 定义输入数据
X = torch.tensor([-1.0, 0.0, 1.0, 2.0, 3.0, 4.0], dtype=torch.float32).unsqueeze(1)
y = torch.tensor([-3.0, -1.0, 1.0, 3.0, 5.0, 7.0], dtype=torch.float32).unsqueeze(1)

# 将输入数据移动到GPU
if torch.cuda.is_available():
    X = X.to('cuda')
    y = y.to('cuda')

# 单层感知机模型
class SingleLayerPerceptron(nn.Module):
    def __init__(self):
        super(SingleLayerPerceptron, self).__init__()
        self.fc = nn.Linear(1, 1)

    def forward(self, x):
        return self.fc(x)

model_single = SingleLayerPerceptron()
if torch.cuda.is_available():
    model_single = model_single.to('cuda')

criterion = nn.MSELoss()
optimizer = optim.SGD(model_single.parameters(), lr=0.01)

for epoch in range(500):
    y_pred = model_single(X)
    loss = criterion(y_pred, y)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 50 == 0:
        print(f'Epoch [{epoch + 1}/500], Loss: {loss.item():.4f}')

with torch.no_grad():
    x_test = torch.tensor([10.0]).unsqueeze(0).unsqueeze(1)
    if torch.cuda.is_available():
        x_test = x_test.to('cuda')
    y_test = model_single(x_test)
    print(f'测试输入x: {x_test.item()}')
    print(f'测试输出y: {y_test.item()}')

# 保存模型
torch.save(model_single.state_dict(), 'precetion_model.path')

# 多层感知机模型
class MultiLayerPerceptron(nn.Module):
    def __init__(self):
        super(MultiLayerPerceptron, self).__init__()
        self.fc1 = nn.Linear(1, 64)
        self.fc2 = nn.Linear(64, 1)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model_multi = MultiLayerPerceptron()
if torch.cuda.is_available():
    model_multi = model_multi.to('cuda')

criterion = nn.MSELoss()
optimizer = optim.SGD(model_multi.parameters(), lr=0.01)

for epoch in range(500):
    y_pred = model_multi(X)
    loss = criterion(y_pred, y)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 50 == 0:
        print(f'Epoch [{epoch + 1}/500], Loss: {loss.item():.4f}')

with torch.no_grad():
    x_test = torch.tensor([10.0]).unsqueeze(0).unsqueeze(1)
    if torch.cuda.is_available():
        x_test = x_test.to('cuda')
    y_test = model_multi(x_test)
    print(f'测试输入x: {x_test.item()}')
    print(f'测试输出y: {y_test.item()}')

# 保存模型
torch.save(model_multi.state_dict(), 'net_model.path')